#include <iostream>
#include <fstream>
#include <vector>
#include <cassert>

namespace IntegrateScore {

    template<typename T>
    T *load_vecs(const char *filename, int &n_data, int &dim) {
        std::ifstream in(filename, std::ios::binary);
        if (!in.is_open()) {
            std::cerr << "Open file error" << std::endl;
            exit(-1);
        }

        in.read((char *) &dim, 4);

        in.seekg(0, std::ios::end);
        std::ios::pos_type ss = in.tellg();
        size_t fsize = (size_t) ss;
        n_data = (int) (fsize / (dim + 1) / 4);

        T *data = new T[(size_t) n_data * (size_t) dim];

        in.seekg(0, std::ios::beg);
        for (int i = 0; i < n_data; i++) {
            in.seekg(4, std::ios::cur);
            in.read((char *) (data + i * dim), dim * sizeof(T));
        }
        in.close();

        return data;
    }

    std::vector<int *> read_label_l_l(char *base_fname, int n_classifier, int &out_n_item, int &out_n_overlap) {
        //label_l_l (n_classifier * n_item * n_overlap)
        std::vector<int *> label_l_l(n_classifier);
        char tmp_fname[256];
        int n_item, n_overlap;
        for (int i = 0; i < n_classifier; i++) {
            std::sprintf(tmp_fname, "%s/Classifier_%d/dataset_partition/label.ivecs", base_fname, i);
            int *label_l = load_vecs<int>(tmp_fname, n_item, n_overlap);
            label_l_l[i] = label_l;
        }
        out_n_item = n_item;
        out_n_overlap = n_overlap;
        return label_l_l;
    }

    std::vector<float *> read_weight_l_l(char *base_fname, int n_classifier, int &out_n_item, int &out_n_overlap) {
        //weight_l_l (n_classifier * n_item * n_overlap)
        std::vector<float *> weight_l_l(n_classifier);
        char tmp_fname[256];
        int n_item, n_overlap;
        for (int i = 0; i < n_classifier; i++) {
            std::sprintf(tmp_fname, "%s/Classifier_%d/dataset_partition/weight_l.fvecs", base_fname, i);
            float *weight_l = load_vecs<float>(tmp_fname, n_item, n_overlap);
            weight_l_l[i] = weight_l;
        }
        out_n_item = n_item;
        out_n_overlap = n_overlap;
        return weight_l_l;
    }

    std::vector<float *>
    read_cluster_score_l(char *base_fname, int n_classifier, int &out_n_query, int &out_n_cluster) {
        //cluster_score_l (n_classifier * n_query * n_cluster)
        std::vector<float *> cluster_score_l(n_classifier);
        char tmp_fname[256];
        int n_query, n_cluster;
        for (int i = 0; i < n_classifier; i++) {
            std::sprintf(tmp_fname, "%s/Classifier_%d/train_eval_model/cluster_score.fvecs", base_fname, i);
            float *cluster_score = load_vecs<float>(tmp_fname, n_query, n_cluster);
            cluster_score_l[i] = cluster_score;
        }
        out_n_query = n_query;
        out_n_cluster = n_cluster;
        return cluster_score_l;
    }

    void save_score_table(char *base_path, std::vector<std::vector<float>> score_table) {
        char res_fname[256];
        std::sprintf(res_fname, "%s/score_table.fvecs", base_path);
        std::ofstream out(res_fname, std::ios::binary | std::ios::out);
        int n_query = (int) score_table.size();
        int n_item = (int) score_table[0].size();
        for (int i = 0; i < n_query; i++) {
            out.write((char *) &n_item, sizeof(unsigned));
            out.write((char *) score_table[i].data(), n_item * sizeof(unsigned));
        }
        out.close();
    }

}
